#!/usr/bin/env python

import os
import struct
import heapq


# just some constants
EvtStreamBegin=0
EvtStreamEnd=1
EvtEnterState=2
EvtLeaveState=3
EvtUpdateCounter=4
EvtSend=5
EvtBroadcast=6
EvtRecv=7

InvalidId = 0xffffffffffffffff
    
class EventRecord:
    def __init__(self, session, ts, et, eid, a, b):
        self.timestamp = ts
        self.type = et
        self.id = eid
        if self.type == EvtStreamBegin:
            if a != InvalidId:
                self.parent = a
            else:
                self.parent = None
            self.tag = b
        elif self.type == EvtStreamEnd:
            self.tag = b
            
        elif self.type == EvtEnterState:
            self.tag = b
        elif self.type == EvtLeaveState:
            self.tag = b

        elif self.type == EvtUpdateCounter:
            self.value = a
            self.tag = b
            
        elif self.type == EvtSend:
            self.length = a
            self.tag = b
        elif self.type == EvtBroadcast:
            self.length = a
            self.tag = b
        elif self.type == EvtRecv:
            self.length = a
            self.tag = b
        else:
            raise "Unknown event"

        def __lt__(self, other):
            return self.timestamp < other.timestamp
        def __le__(self, other):
            return self.timestamp <= other.timestamp


class State:
    def __init__(self, session, sid, name):
        self.session = session
        self.id = sid
        self.name = name

    def __str__(self):
        return str(self.name)

class Counter:
    def __init__(self, session, cid, name):
        self.session = session
        self.id = cid
        self.name = name

    def __str__(self):
        return str(self.name)

class StreamEvent:
    def getParentName(self):
        if self.parent is not None:
            if self.parent.__dict__.has_key("name"):
                return self.parent.name
            if self.parent.__dict__.has_key("state"):
                return self.parent.state.name
        return ''

class StateChangeEvent(StreamEvent):
    def __init__(self, state, start, tag, parent, stream):
        self.state = state
        self.start = start
        self.starttag = tag
        self.finish = None
        self.finishtag = None
        self.children = []
        self.parent = parent
        self.stream = stream
        self.duration = 0
        
        if parent is None:
            self.depth = 0
            stream.addChild(self)
        else:
            self.depth = parent.depth + 1
            parent.addChild(self)

    def addChild(self, x):
        x.parent = self
        self.children.append(x)

    def __repr__(self):
        return "StateChangeEvent(id={sid}, name='{name}', start={start}, starttag={starttag}, finish={finish}, finishtag={finishtag}, stream='{stream}', parent='{parent}', children={children})".format(
                sid=self.state.id,
                name=self.state.name,
                start=self.start,
                starttag=self.starttag,
                finish=self.finish,
                finishtag=self.finishtag,
                stream=self.stream.name,
                parent=self.getParentName(),
                children=len(self.children))

    def pprint(self, indent=""):
        print "{time:016}{indent} +> {name} tag={tag}".format(time=self.start, indent=indent, name=self.state.name, tag=self.starttag)
        for c in self.children:
            c.pprint(indent+" | ")
        print "{time:016}{indent} +< {name} tag={tag}".format(time=self.finish, indent=indent, name=self.state.name, tag=self.finishtag)


class Receiver:
    def __init__(self, timestamp, state, length, tag):
        self.time = timestamp
        self.state = state
        self.length = length
        self.tag = tag



class CounterUpdateEvent(StreamEvent):
    def __init__(self, timestamp, counter, value, tag, parent, stream):
        self.time = timestamp
        self.counter = counter
        self.value = value
        self.tag = tag
        self.parent = parent
        self.stream = stream
        if parent is None:
            self.depth = 0
            stream.addChild(self)
        else:
            self.depth = parent.depth + 1
            parent.addChild(self)

    def __repr__(self):
        return "CounterUpdateEvent(id={sid}, name='{name}', value={value}, tag={tag}, stream='{stream}', parent='{parent}')".format(
                sid=self.counter.id,
                name=self.counter.name,
                value=self.value,
                tag=self.tag,
                stream=self.stream.name,
                parent=self.getParentName())

    def pprint(self, indent):
        print "{time:016}{indent} == {name}={value} tag={tag}".format(time=self.time, indent=indent, name=self.counter.name, value=self.value, tag=self.tag)


class Stream:
    def __init__(self, session, sid, name):
        self.id = sid
        self.name = name
        self.start = None
        self.finish = None
        self.parent = None
        self.children = []
        self.session = session
        self.starttag = None
        self.finishtag = None
        self.duration = None

        f = file(session.filename + "-" + self.name + ".stream", "rb")
        evsize=40
        header = f.read(evsize)
        streamMagic = "KOTF-1.0.0stream";
        if header.startswith(streamMagic) == False:
            raise "wrong header!"

        ststack = [None] # there is always None at the bottom
        
        while True:
            buf = f.read(evsize)
            if len(buf) != evsize:
                break
            
            (timestamp,etype,eid,a,b) = struct.unpack("<QIxxxxQQQ", buf)
            e = EventRecord(self.session, timestamp, etype, eid, a, b)

            if e.type == EvtStreamBegin:
                
                self.start = e.timestamp
                self.starttag = e.tag
                self.parentid = e.parent

            elif e.type == EvtStreamEnd:

                self.finish = e.timestamp
                self.finishtag = e.tag
                self.duration = self.finish - self.start

            elif e.type == EvtEnterState:

                parent = ststack[-1]
                st = StateChangeEvent(self.session.states[e.id], e.timestamp, e.tag, parent, self)
                ststack.append(st)
                
            elif e.type == EvtLeaveState:

                st = ststack.pop()
                assert e.id == st.state.id
                st.finish = e.timestamp
                st.finishtag = e.tag
                st.duration = st.finish - st.start

            elif e.type == EvtUpdateCounter:
                
                parent = ststack[-1]
                cnt = CounterUpdateEvent(e.timestamp, self.session.counters[e.id], e.value, e.tag, parent, self)

    def addChild(self, st):
        self.children.append(st)

    def getParentName(self):
        if self.parent is not None:
            return self.parent.name
        else:
            return ''

    def __repr__(self):
        return "Stream(name='%s', start=%u, starttag=%u, finish=%u, finishtag=%u, parent='%s', children=%s)" % \
               (self.name, self.start, self.starttag, self.finish, self.finishtag, self.getParentName(), repr(self.children))

    def pprint(self, indent=""):
        print "Stream '%s' started at %u tag=%u parent='%s'" % (self.name, self.start, self.starttag, self.getParentName())
        for s in self.children:
            s.pprint(indent+" ")
        print "Stream '%s' finished at %u tag=%u" % (self.name, self.finish, self.finishtag)




class Session:

    def __init__(self, sessfilename):
        if sessfilename[-5:] != ".kotf":
            sessfilename += ".kotf"

        self.filename = sessfilename[:-5]
        self.name = os.path.basename(self.filename)

        self.start = 0
        self.finish = 0
        self.duration = 0
        self.streamNames={}
        self.stateNames={}
        self.counterNames={}
        self.channelNames={}

        ld =  {
            "startTime": self.start,
            "finishTime": self.finish
            }
        execfile(sessfilename, {
            "stream": self.streamNames,
            "state": self.stateNames,
            "counter": self.counterNames,
            "channel": self.channelNames
            }, ld)
        self.start = ld["startTime"]
        self.finish = ld["finishTime"]
        self.duration - self.finish - self.start

        self.states = {}
        for sid,statename in self.stateNames.iteritems():
            self.states[sid] = State(self, sid, statename)

        self.counters = {}
        for cid, cname in self.counterNames.iteritems():
            self.counters[cid] = Counter(self, cid, cname)

        self.streams = [None] * (len(self.streamNames) - 1) # HACK
        for sid,streamname in self.streamNames.iteritems():
            if streamname == "scheduler":
                continue # HACK
            self.streams[sid] = Stream(self, sid, streamname)
        del self.streamNames

        for s in self.streams:
            if s.parentid is not None:
                s.parent = self.streams[s.parentid]

    def __repr__(self):
        return "Session(name='%s', start=%u, finish=%u, streams=%s)" % \
            (self.name, self.start, self.finish, repr(self.streams))

    def pprint(self):
        print "Session '%s' started at %u" % (self.name, self.start)
        for s in self.streams:
            s.pprint("")
        print "Session '%s' finished at %u" % (self.name, self.finish)



if __name__ == "__main__":
    import sys
    if len(sys.argv) != 2:
        print "usage: ./kotf.py <sessionfile.kotf>"
        sys.exit(1)
    session = Session(sys.argv[1])
    session.pprint()
